#include "klib.h"
#include "vme.h"
#include "cte.h"
#include "loader.h"
#include "disk.h"
#include "fs.h"
#include <elf.h>

uint32_t load_elf(PD *pgdir, const char *name) {
  Elf32_Ehdr elf;
  Elf32_Phdr ph;
  inode_t *inode = iopen(name, TYPE_NONE);
  if (!inode) return -1;
  iread(inode, 0, &elf, sizeof(elf));
  if (*(uint32_t*)(&elf) != 0x464c457f) { // check ELF magic number
    iclose(inode);
    return -1;
  }
  for (int i = 0; i < elf.e_phnum; ++i) {
    iread(inode, elf.e_phoff + i * sizeof(ph), &ph, sizeof(ph));
    if (ph.p_type == PT_LOAD) {
      // Lab1-2: Load segment to physical memory
      // Lab1-4: Load segment to virtual memory
      // void *dst1 = (void*)(ph.p_vaddr);
      // uint32_t filesz_p = ph.p_filesz;
      // iread(inode, ph.p_offset, dst1, filesz_p);
      
      // void *dst2 = (void*)(ph.p_vaddr + ph.p_filesz);
      // uint32_t filesz_0 = ph.p_memsz - ph.p_filesz;
      // memset(dst2, 0, filesz_0);
      
      //记录原来的pgdir
      PD *ori_pgdir =  vm_curr();
      //切换到新的pgdir
      set_cr3(pgdir);

    //将这一段所需要的虚拟内存（即[p_vaddr, p_vaddr+p_memsz)）添加到pgdir管理的映射，你可以使用前面实现的vm_map函数
    //todo!!!!这里的prot应该设置为多少？
    // 测试点
    // printf("ph.p_vaddr: %x\n", ph.p_vaddr);
    vm_map(pgdir, ph.p_vaddr, ph.p_memsz, 7);




    //将这一段的内容从文件中读取到内存中（说是是直接利用qemu就可以访问虚拟内存）
    void *dst1 = (void*)(ph.p_vaddr);
    uint32_t filesz_p = ph.p_filesz;
    iread(inode, ph.p_offset, dst1, filesz_p);
      
    void *dst2 = (void*)(ph.p_vaddr + ph.p_filesz);
    uint32_t filesz_0 = ph.p_memsz - ph.p_filesz;
    memset(dst2, 0, filesz_0);







      //Load segment to virtual memory    AI生成的，不一定有用
      // for (uint32_t va = ROUNDDOWN(ph.p_vaddr, PGSIZE); va < ph.p_vaddr + ph.p_memsz; va += PGSIZE) {
      //   PTE *pte = vm_walkpgdir(pgdir, va, 1);
      //   if (pte == NULL) {
      //     iclose(inode);
      //     return -1;
      //   }
      //   if (pte->present == 0) {
      //     uint32_t pa = kalloc();
      //     if (pa == 0) {
      //       iclose(inode);
      //       return -1;
      //     }
      //     pte->val = pa | PTE_P | PTE_W | PTE_U;
      //   }
      // }



      //恢复原来的pgdir
      set_cr3(ori_pgdir);
    }
  }
  // TODO: Lab1-4 alloc stack memory in pgdir
  //给pgdir映射上[0xbffff000, 0xc0000000)（即[USR_MEM-PGSIZE, USR_MEM)）这一块虚拟内存作为用户程序的栈。
  vm_map(pgdir, USR_MEM - PGSIZE, PGSIZE, 7);



  iclose(inode);
  return elf.e_entry;
}

#define MAX_ARGS_NUM 31

uint32_t load_arg(PD *pgdir, char *const argv[]) {
  // Lab1-8: Load argv to user stack
  char *stack_top = (char*)vm_walk(pgdir, USR_MEM - PGSIZE, 7) + PGSIZE;
  size_t argv_va[MAX_ARGS_NUM + 1];
  int argc;
  for (argc = 0; argv[argc]; ++argc) {
    assert(argc < MAX_ARGS_NUM);
    // push the string of argv[argc] to stack, record its va to argv_va[argc]
    size_t len = strlen(argv[argc]) + 1;
    stack_top -= len;
    strcpy(stack_top, argv[argc]);
    // str_cpy((char*)vm_walk(pgdir, ADDR2PAGE(stack_top), 7), argv[argc]);
    // memcpy(stack_top, argv[argc], len);
    argv_va[argc] = USR_MEM - PGSIZE + ADDR2OFF(stack_top);
    // argv_va[argc] = USR_MEM - PGSIZE + ADDR2OFF(stack_top);

    
  }
  argv_va[argc] = 0; // set last argv NULL
  stack_top -= ADDR2OFF(stack_top) % 4; // align to 4 bytes
  for (int i = argc; i >= 0; --i) {
    // push the address of argv_va[argc] to stack to make argv array
    stack_top -= sizeof(size_t);
    *(size_t*)stack_top = argv_va[i];
  }
  // push the address of the argv array as argument for _start
  size_t argv0_addr = USR_MEM - PGSIZE + ADDR2OFF(stack_top);

  stack_top -= sizeof(size_t);
  *(size_t*)stack_top = argv0_addr;

  // push argc as argument for _start
  stack_top -= sizeof(size_t);
  *(size_t*)stack_top = argc;
  stack_top -= sizeof(size_t); // a hole for return value (useless but necessary)
  return USR_MEM - PGSIZE + ADDR2OFF(stack_top);
}

int load_user(PD *pgdir, Context *ctx, const char *name, char *const argv[]) {
  size_t eip = load_elf(pgdir, name);
  if (eip == -1) return -1;
  ctx->cs = USEL(SEG_UCODE);
  ctx->ds = USEL(SEG_UDATA);
  ctx->eip = eip;
  // TODO: Lab1-6 init ctx->ss and esp
  ctx->ss = USEL(SEG_UDATA);
  
  ctx->esp = load_arg(pgdir, argv);

  /*之前的esp
  ctx->esp = USR_MEM-16;
*/

  ctx->eflags = 0x202; // TODO: Lab1-7 change me to 0x202
  return 0;
}
